import tensorflow as tf
from tensorflow.python.keras.metrics import MeanMetricWrapper


def nmse(x, x_hat):
    b = tf.shape(x)[0]

    x = tf.reshape(x, [b, -1, 2])
    x_real, x_imag = x[:, :, 0], x[:, :, 1]

    x_hat = tf.reshape(x_hat, [b, -1, 2])
    x_hat_real, x_hat_imag = x_hat[:, :, 0], x_hat[:, :, 1]

    x_real = tf.cast(x_real, tf.complex64)
    x_imag = tf.cast(x_imag, tf.complex64)
    x_hat_real = tf.cast(x_hat_real, tf.complex64)
    x_hat_imag = tf.cast(x_hat_imag, tf.complex64)

    x_C = x_real - 0.5 + 1j * (x_imag - 0.5)
    x_hat_C = x_hat_real - 0.5 + 1j * (x_hat_imag - 0.5)
    power = tf.reduce_sum(tf.abs(x_C) ** 2, axis=1)
    mse = tf.reduce_sum(tf.abs(x_C - x_hat_C) ** 2, axis=1)
    result = mse / power
    return result


class NMSE(MeanMetricWrapper):

    def __init__(self, name='nmse', dtype=None):
        super().__init__(nmse, name, dtype=dtype)
